import torch
import torch.nn as nn
import torch.optim as optim
import time
import numpy as np


# 方法1: 高斯消元
def solve_gaussian_elimination(K, F):
    # 高斯消元通常需要在 CPU 上实现
    K_cpu = K.cpu().numpy()
    F_cpu = F.cpu().numpy()
    solution = np.linalg.solve(K_cpu, F_cpu)
    return torch.tensor(solution, dtype=K.dtype, device=K.device)


# 方法2: 使用SVD
def solve_svd(K, F):
    U, S, Vh = torch.linalg.svd(K)
    S_inv = torch.diag(1.0 / S)
    K_svd_inv = Vh.T @ S_inv @ U.T
    return K_svd_inv @ F


# 方法3: 使用LU分解
def solve_lu(K, F):
    P, L, U = torch.linalg.lu(K)
    y = torch.linalg.solve(L, F)
    u = torch.linalg.solve(U, y)
    return u


# 定义一个简单的神经网络
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(3, 10)  # 输入3，表示三种方法
        self.fc2 = nn.Linear(10, 3)  # 输出3，表示三种方法的概率
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return self.softmax(x)


# 生成随机数据
def generate_data(n):
    K = torch.randn(n, n)
    F = torch.randn(n)
    return K, F


# 训练模型
def train_model(model, data, labels, epochs=10):
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    for epoch in range(epochs):
        for i in range(len(data)):
            optimizer.zero_grad()
            outputs = model(data[i].unsqueeze(0))
            loss = criterion(outputs, labels[i].unsqueeze(0))
            loss.backward()
            optimizer.step()


# 主程序
if __name__ == "__main__":
    # 配置
    num_samples = 100000
    max_n = 20
    model = SimpleNN()

    # 生成数据和标签
    data = []
    labels = []
    for _ in range(num_samples):
        n = np.random.randint(1, max_n + 1)
        K, F = generate_data(n)

        # 记录时间
        start = time.time()
        u1 = solve_gaussian_elimination(K, F)
        time_gaussian = time.time() - start

        start = time.time()
        u2 = solve_svd(K, F)
        time_svd = time.time() - start

        start = time.time()
        u3 = solve_lu(K, F)
        time_lu = time.time() - start

        # 选择最快的方法
        times = [time_gaussian, time_svd, time_lu]
        fastest_method = times.index(min(times))

        # 将数据和标签添加到列表
        data.append(torch.tensor(times, dtype=torch.float32))
        labels.append(torch.tensor(fastest_method, dtype=torch.long))

    # 转换为Tensor
    data = torch.stack(data)
    labels = torch.stack(labels)

    # 训练模型
    train_model(model, data, labels)

    # 测试模型
    n_test = np.random.randint(1, max_n + 1)
    K_test, F_test = generate_data(n_test)

    # 记录时间
    start = time.time()
    solve_gaussian_elimination(K_test, F_test)
    time_gaussian_test = time.time() - start

    start = time.time()
    solve_svd(K_test, F_test)
    time_svd_test = time.time() - start

    start = time.time()
    solve_lu(K_test, F_test)
    time_lu_test = time.time() - start

    # 生成测试数据
    test_data = torch.tensor(
        [time_gaussian_test, time_svd_test, time_lu_test], dtype=torch.float32
    )
    model.eval()
    with torch.no_grad():
        output = model(test_data.unsqueeze(0))
        predicted_method = torch.argmax(output, dim=-1).item()

    print(f"Predicted fastest method: Method {predicted_method}")
